#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
验证增强版LSTM-CRF模型预测修复是否有效
"""

import os
import sys
import torch

# 添加项目根目录到路径
project_root = os.path.dirname(os.path.abspath(__file__))
sys.path.append(project_root)

def verify_prediction_fix():
    """验证预测修复"""
    try:
        print("验证增强版LSTM-CRF模型预测修复...")
        
        # 导入必要的模块
        from algorithms.enhanced_lstm_crf import EnhancedLstmCRFTrainer
        
        # 创建训练器实例来获取正确的配置
        trainer = EnhancedLstmCRFTrainer(
            csv_file_path=os.path.join(project_root, 'scripts', 'plw', 'plw_history.csv'),
            model_save_path=os.path.join(project_root, 'scripts', 'plw', 'enhanced_lstm_crf_model.pth'),
            lottery_type='plw'
        )
        
        print(f"训练器配置 - input_dim: {trainer.input_dim}")
        print(f"训练器配置 - output_dim: {trainer.output_dim}")
        print(f"训练器配置 - output_seq_length: {trainer.output_seq_length}")
        
        # 加载模型
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        model_path = os.path.join(project_root, 'scripts', 'plw', 'enhanced_lstm_crf_model.pth')
        
        if os.path.exists(model_path):
            checkpoint = torch.load(model_path, map_location=device)
            print(f"模型检查点键: {checkpoint.keys()}")
            
            if 'model_config' in checkpoint:
                config = checkpoint['model_config']
                print(f"模型配置: {config}")
                
                # 检查输入维度是否正确
                if config.get('input_dim') == 10:
                    print("✅ 模型配置中的输入维度正确 (10)")
                else:
                    print(f"❌ 模型配置中的输入维度不正确: {config.get('input_dim')}")
                    
        else:
            print(f"模型文件不存在: {model_path}")
            
    except Exception as e:
        print(f"验证过程中出错: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    verify_prediction_fix()

# 测试process_predictions函数
def test_process_predictions():
    """测试process_predictions函数对排列5的处理"""
    
    def process_predictions(red_predictions, blue_predictions, lottery_type):
        """
        处理预测结果，确保号码在有效范围内且为整数
        :param red_predictions: list, 红球预测的类别索引（排列5时为全部数字）
        :param blue_predictions: list, 蓝球预测的类别索引（排列5时为None）
        :param lottery_type: 'ssq'、'dlt' 或 'plw'
        :return: list, 预测的开奖号码
        """
        if lottery_type == "dlt":
            # 大乐透前区：1-35，后区：1-12
            front_numbers = [min(max(int(num) + 1, 1), 35) for num in red_predictions[:5]]
            back_numbers = [min(max(int(num) + 1, 1), 12) for num in blue_predictions[:2]]

            # 确保前区号码唯一
            front_numbers = list(set(front_numbers))
            while len(front_numbers) < 5:
                additional_num = 1  # 简化处理
                if additional_num not in front_numbers:
                    front_numbers.append(additional_num)
            front_numbers = sorted(front_numbers)[:5]
            
            return front_numbers + back_numbers

        elif lottery_type == "ssq":
            # 双色球红球：1-33，蓝球：1-16
            front_numbers = [min(max(int(num) + 1, 1), 33) for num in red_predictions[:6]]
            back_number = min(max(int(blue_predictions[0]) + 1, 1), 16)

            # 确保红球号码唯一
            front_numbers = list(set(front_numbers))
            while len(front_numbers) < 6:
                additional_num = 1  # 简化处理
                if additional_num not in front_numbers:
                    front_numbers.append(additional_num)
            front_numbers = sorted(front_numbers)[:6]
            
            return front_numbers + [back_number]
            
        elif lottery_type == "plw":
            # 排列5：5个数字，每个数字范围 0-9
            # red_predictions 实际上包含了5个数字位置的预测
            numbers = [min(max(int(num), 0), 9) for num in red_predictions[:5]]
            return numbers  # 返回处理后的数字列表

    print("🔧 测试process_predictions函数对排列5的处理")
    print("=" * 60)
    
    # 测试排列5预测结果处理
    try:
        # 模拟模型预测结果
        red_predictions = [1, 2, 3, 4, 5]
        blue_predictions = None  # 排列5没有蓝球预测
        
        result = process_predictions(red_predictions, blue_predictions, "plw")
        
        print(f"✅ 排列5预测结果处理成功")
        print(f"   输入: red_predictions={red_predictions}, blue_predictions={blue_predictions}")
        print(f"   输出: {result}")
        print(f"   类型: {type(result)}")
        
        # 验证结果格式
        if isinstance(result, list) and len(result) == 5:
            print(" 返回结果格式正确")
            return True
        else:
            print(" 返回结果格式不正确")
            return False
            
    except Exception as e:
        print(f" 排列5预测结果处理失败: {e}")
        return False

# 测试run_multi_algorithm_prediction中的结果处理
def test_multi_algorithm_result_processing():
    """测试多算法预测结果处理"""
    print("\n🔧 测试多算法预测结果处理")
    print("=" * 60)
    
    # 模拟算法预测结果
    algo_predictions = [
        {
            'red': [1, 2, 3, 4, 5],  # 排列5的5个数字
            'blue': []  # 排列5没有蓝球
        }
    ]
    
    algorithm_name = "🤖 LSTM-CRF 经典模式"
    lottery_type = "plw"
    
    predictions = []
    
    try:
        # 模拟run_multi_algorithm_prediction中的处理逻辑
        for pred in algo_predictions:
            # 确保pred是字典且包含必要的键
            if isinstance(pred, dict):
                # 获取预测数字，确保不是None
                red_balls = pred.get('red', [])
                blue_balls = pred.get('blue', [])
                
                # 对于排列5，如果red_balls是None，则使用空列表
                if red_balls is None:
                    red_balls = []
                
                predictions.append({
                    'algorithm': algorithm_name,
                    'confidence': pred.get('confidence', 0.8),  # 模拟置信度
                    'red': red_balls,
                    'blue': blue_balls
                })
            else:
                # 如果pred不是字典，创建一个默认格式
                # 根据彩票类型设置默认值
                if lottery_type == "plw":
                    # 排列5：5个数字，蓝球为空
                    default_red = []
                    default_blue = []
                else:
                    # 大乐透和双色球使用空列表
                    default_red = []
                    default_blue = []
                
                predictions.append({
                    'algorithm': algorithm_name,
                    'confidence': 0.8,  # 模拟置信度
                    'red': default_red,
                    'blue': default_blue
                })
        
        print(f"✅ 多算法预测结果处理成功")
        print(f"   处理后的预测结果: {predictions}")
        
        # 验证结果格式
        if (isinstance(predictions, list) and 
            len(predictions) > 0 and 
            isinstance(predictions[0], dict) and
            'red' in predictions[0] and
            'blue' in predictions[0]):
            
            red_balls = predictions[0]['red']
            blue_balls = predictions[0]['blue']
            
            print(f"   红球数字: {red_balls} (类型: {type(red_balls)})")
            print(f"   蓝球数字: {blue_balls} (类型: {type(blue_balls)})")
            
            # 对于排列5，红球应该是包含5个数字的列表，蓝球应该是空列表
            if (isinstance(red_balls, list) and 
                isinstance(blue_balls, list) and
                len(red_balls) == 5 and
                len(blue_balls) == 0):
                print(" 预测结果格式正确")
                return True
            else:
                print(" 预测结果格式不正确")
                return False
        else:
            print(" 预测结果格式不正确")
            return False
            
    except Exception as e:
        print(f" 多算法预测结果处理失败: {e}")
        return False

if __name__ == "__main__":
    print("🎯 验证排列5多算法预测修复")
    print("=" * 80)
    
    # 运行所有测试
    test1 = test_process_predictions()
    test2 = test_multi_algorithm_result_processing()
    
    print("\n" + "=" * 80)
    if all([test1, test2]):
        print("🎉 所有修复验证通过!")
        print(" process_predictions函数修复")
        print(" 多算法预测结果处理修复")
    else:
        print("💥 部分修复验证失败!")
        if not test1:
            print(" process_predictions函数修复验证失败")
        if not test2:
            print(" 多算法预测结果处理验证失败")